#!/usr/bin/env python3
# Author: Armit
# Create Time: 2024/09/07 

# 查看高斯模型的 ply 文件属性的统计信息

from pathlib import Path
from argparse import ArgumentParser
import pandas as pd
from ydata_profiling import ProfileReport

import numpy as np
from numpy import ndarray
from plyfile import PlyData, PlyElement
import matplotlib.pyplot as plt


def stats_array(x:ndarray, name:str):
  print(f'[{name}]')
  print(f'  max: {np.max(x):.4f}')
  print(f'  min: {np.min(x):.4f}')
  print(f'  avg: {np.mean(x):.4f}')
  print(f'  std: {np.std(x):.4f}')

def run(args):
  plydata = PlyData.read(args.fp)
  vertices: PlyElement = plydata['vertex']
  props = [p for p in vertices.properties if p.name not in ['nx', 'ny', 'nz']]    # ignore all zero normals
  print(f'>> file: {args.fp}')
  print(f'>> n_pts: {len(vertices)}')
  print(f'>> props: {[p.name for p in props]}')

  # pandas profiler
  profile_fp = Path(args.fp).with_suffix('.html')
  if not profile_fp.exists():
    keys = [p.name for p in props if 'f_rest_' not in p.name]   # ignore higher HS coeffs
    vals = [vertices[k] for k in keys]
    df = pd.DataFrame({k: v for k, v in zip(keys, vals)})
    print('>> making profile report...')
    ProfileReport(df, explorative=True).to_file(profile_fp)
    print(f'>> save to {profile_fp}')

  FILE_FOLDER = Path(args.fp).parent

  # xyz
  x = vertices['x'] ; stats_array(x, 'x')
  y = vertices['y'] ; stats_array(y, 'y')
  z = vertices['z'] ; stats_array(z, 'z')

  plt.clf()
  plt.figure(figsize=(8, 8))
  plt.subplot(221) ; plt.hist(x, bins=100) ; plt.title('x')
  plt.subplot(222) ; plt.hist(y, bins=100) ; plt.title('y')
  plt.subplot(223) ; plt.hist(z, bins=100) ; plt.title('z')
  plt.suptitle('xyz')
  plt.tight_layout()
  fp = FILE_FOLDER / 'xyz.png'
  plt.savefig(fp, dpi=600)
  print(f'>> savefig {fp}')

  # opacity, scale
  scale_0 = vertices['scale_0'] ; stats_array(scale_0, 'scale_0')
  scale_1 = vertices['scale_1'] ; stats_array(scale_1, 'scale_1')
  scale_2 = vertices['scale_2'] ; stats_array(scale_2, 'scale_2')
  opacity = vertices['opacity'] ; stats_array(opacity, 'opacity')

  plt.clf()
  plt.figure(figsize=(8, 8))
  plt.subplot(221) ; plt.hist(scale_0, bins=100) ; plt.title('scale_0')
  plt.subplot(222) ; plt.hist(scale_1, bins=100) ; plt.title('scale_1')
  plt.subplot(223) ; plt.hist(scale_2, bins=100) ; plt.title('scale_2')
  plt.subplot(224) ; plt.hist(opacity, bins=100) ; plt.title('opacity')
  plt.suptitle('scale + opacity')
  plt.tight_layout()
  fp = FILE_FOLDER / 'scale_opacity.png'
  plt.savefig(fp, dpi=600)
  print(f'>> savefig {fp}')

  # rot
  rot_0 = vertices['rot_0'] ; stats_array(rot_0, 'rot_0')
  rot_1 = vertices['rot_1'] ; stats_array(rot_1, 'rot_1')
  rot_2 = vertices['rot_2'] ; stats_array(rot_2, 'rot_2')
  rot_3 = vertices['rot_3'] ; stats_array(rot_3, 'rot_3')

  plt.clf()
  plt.figure(figsize=(8, 8))
  plt.subplot(221) ; plt.hist(rot_0, bins=100) ; plt.title('rot_0')
  plt.subplot(222) ; plt.hist(rot_1, bins=100) ; plt.title('rot_1')
  plt.subplot(223) ; plt.hist(rot_2, bins=100) ; plt.title('rot_2')
  plt.subplot(224) ; plt.hist(rot_3, bins=100) ; plt.title('rot_3')
  plt.suptitle('rot')
  plt.tight_layout()
  fp = FILE_FOLDER / 'rot.png'
  plt.savefig(fp, dpi=600)
  print(f'>> savefig {fp}')


if __name__ == '__main__':
  parser = ArgumentParser()
  parser.add_argument('-F', '--fp', required=True, help='path to *.ply file')
  args = parser.parse_args()

  run(args)
